DFF_BIN = r"D:\dff7b-Generic\dff7b\bin32w"
PPF_FILE = r"TEAM20.ppf"
DATA_FOLDER = r"D:\DATA\DffForceField\T_TEAM_ver1\intra\Data"
WORK_FOLDER = r"D:\DATA\DffForceField\T_TEAM_ver1\intra\FitAMBER"

FIT0_DFI = """$SYSTEM 
  ROOT = %(DFF_ROOT)s
$END 
$DATA 
    %(QMD_FILE)s
$END 
$MODEL
    %(MSD_FILE)s
$END 
$BUNDLES 
  CONFORMER = QMD
$END 
$FORCEFIELD 
  TYPE = AMBER
  INPUT = %(INPUT_PPF)s
  OUTPUT = %(OUTPUT_PPF)s
  SAVEALL = NO
  EQUI = INPUT
$END 
$PARAMETERS 
  TORSION = DEFAULT 
  WILDCARD = NO 
  NONBOND = NO 
  OOPA = YES 
  UNHARMONIC = NO 
$END 
$COMPUTE  
  FIT = NONL
  SETTypes = YES
  SETCharges = YES
$END 
$FITOPTIONS
  WEIGHTS = 100.0 10.0 1.0
  ITERATION = 0
  CONVERGENCES = 0.001  0.00001
  AUTOCORRECT = YES
  VARIANCES = 5.0
  PENALTY = 100
$END
$PRINT
$END
"""

import os
import kuai

ff, eqt = kuai.read_ff(PPF_FILE)

DFFFIT_COMMAND = os.path.normcase(os.path.join(DFF_BIN, 'dfffit'))
DFFFIT_COMMAND = '"' + DFFFIT_COMMAND + '" '  

def scan_file(qmdfile, msdfile, dfifile):
    global DFF_BIN
    
    workdir = os.path.dirname(dfifile)
    workname = os.path.basename(dfifile)[:-4]
    
    DFF_ROOT = os.path.dirname(DFF_BIN)
    QMD_FILE = qmdfile
    MSD_FILE = msdfile
    INPUT_PPF = PPF_FILE
    OUTPUT_PPF = os.path.join(workdir, workname+".ppf")
    
    dfifile = open(dfifile, "w")
    dfifile.write(FIT0_DFI % vars())
    dfifile.close()
    
    mol = kuai.read_mol(msdfile)
    kuai.assign_charges(mol, ff, eqt)
    kuai.write_mol(msdfile, mol)
    os.chdir(workdir)
    os.system(DFFFIT_COMMAND + workname)
    
def scan_data(folder_data, folder_work):
    subs = os.listdir(folder_data)
    for i in subs:
        if i == "Later":
            continue
        fullname_data = os.path.join(folder_data, i)
        fullname_work = os.path.join(folder_work, i)
        if os.path.isdir(fullname_data):
            if not os.path.isdir(fullname_work):
                os.makedirs(fullname_work)
            scan_data(fullname_data, fullname_work)
        elif i.endswith(".qmd"):
            qmdfile = fullname_data 
            msdfile = qmdfile[:-4] + ".msd"
            dfifile = os.path.join(folder_work, i[:-4] + ".dfi")
            if os.path.isfile(msdfile):
                scan_file(qmdfile, msdfile, dfifile)


if __name__ == '__main__':
    scan_data(DATA_FOLDER, WORK_FOLDER)
